import inspect
import os
from functools import wraps

from fastmcp import FastMCP, Context

app = FastMCP(name="caculator")

def login_required():
    def wrapper(func):
        @wraps(func)
        async def wrapped(*args, **kwargs):
            ctx = app.get_context() # ? 这种方式也是mcp server服务端获取上下文的方式
            request = ctx.get_http_request()
            print(request.headers)
            if inspect.iscoroutinefunction(func): # 是否为异步请求
                return await func(*args, **kwargs)
            else:
                return func(*args, **kwargs)
        return wrapped
    return wrapper

# 装饰器本质上执行的是内部的wrapped函数
# wrapper = login_required()
# wrapper(plus_tool)

# 注意装饰器顺序：@app.tool() 放在上方，@login_required() 放在下方，装饰器的执行顺序是从上到下
@app.tool()
@login_required()
def plus_tool(a: float, b: float, ctx: Context) -> float:
    """
    计算两数相加的结果
    :param a: 相加的第一个数
    :param b: 相加的第二个数
    :return: 返回两数相加的结果
    """
    # ? 以下是在tool函数内部添加ctx参数获取上下文
    # request = ctx.get_http_request()
    # headers = request.headers
    # print(headers)
    return a + b

if __name__ == '__main__':
    app.run(transport='sse')